#include <NTDDK.h>

#define LOOPBACK_DEVICE_NAME L"\\Device\\LOOPBACK"
#define LOOPBACK_SYMLINK_NAME L"\\DosDevices\\LBK"

typedef struct _DEVICE_EXTENSION {
	PDEVICE_OBJECT device;
	PVOID buffer;
	ULONG buffer_size;
} DEVICE_EXTENSION, *PDEVICE_EXTENSION;


static NTSTATUS CreateDevice (PDRIVER_OBJECT driver, ULONG device_number) 
{
	NTSTATUS status;
	PDEVICE_OBJECT device;
	PDEVICE_EXTENSION device_extension;
	UNICODE_STRING device_name;
	UNICODE_STRING symlink_name;

	RtlInitUnicodeString(&device_name, LOOPBACK_DEVICE_NAME);
	RtlInitUnicodeString(&symlink_name, LOOPBACK_SYMLINK_NAME); 

	status = IoCreateDevice(driver,
							sizeof(DEVICE_EXTENSION),
							&device_name,
							FILE_DEVICE_UNKNOWN,
							0, 
							FALSE,
							&device);
						
	if (!NT_SUCCESS(status))
		return status;

	device->Flags |= DO_BUFFERED_IO;

	device_extension = (PDEVICE_EXTENSION)device->DeviceExtension;
	device_extension->device = device;
	device_extension->buffer = NULL;
	device_extension->buffer_size = 0;

	status = IoCreateSymbolicLink(&symlink_name, &device_name);
	if (!NT_SUCCESS(status)) {
		IoDeleteDevice(device);
		return status;
	}

	return STATUS_SUCCESS;
}

static VOID DriverUnload (PDRIVER_OBJECT driver) 
{
	UNICODE_STRING symlink_name;
	PDEVICE_OBJECT next_device = driver->DeviceObject;
	RtlInitUnicodeString(&symlink_name, LOOPBACK_SYMLINK_NAME);

	while (next_device != NULL) {
		PDEVICE_EXTENSION device_extension = (PDEVICE_EXTENSION) next_device->DeviceExtension;
		if (device_extension->buffer != NULL) {
			ExFreePool(device_extension->buffer);
			device_extension->buffer = NULL;
			device_extension->buffer_size = 0;
		}
		 
		IoDeleteSymbolicLink(&symlink_name);
		next_device = next_device->NextDevice;
		IoDeleteDevice(device_extension->device);
	}
	return;
}

static NTSTATUS Create (PDEVICE_OBJECT	device, PIRP irp) 
{
	irp->IoStatus.Status = STATUS_SUCCESS;
	irp->IoStatus.Information = 0;	
	IoCompleteRequest(irp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}

static NTSTATUS Close (IN PDEVICE_OBJECT device, IN PIRP irp)
{
	PDEVICE_EXTENSION device_extension = (PDEVICE_EXTENSION) device->DeviceExtension;

	if (device_extension->buffer != NULL) {
		ExFreePool(device_extension->buffer);
		device_extension->buffer = NULL;
		device_extension->buffer_size = 0;
	}

	irp->IoStatus.Status = STATUS_SUCCESS;
	irp->IoStatus.Information = 0;
	IoCompleteRequest(irp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}


static NTSTATUS Write (IN PDEVICE_OBJECT device, IN PIRP irp) 
{	
	NTSTATUS status = STATUS_SUCCESS;
	PVOID buffer;
	ULONG size;

	PIO_STACK_LOCATION irpStack = IoGetCurrentIrpStackLocation(irp);
	PDEVICE_EXTENSION device_extension = (PDEVICE_EXTENSION)device->DeviceExtension;

	if (device_extension->buffer != NULL) {
		ExFreePool(device_extension->buffer);
		device_extension->buffer = NULL;
		device_extension->buffer_size = 0;
	}

	size = irpStack->Parameters.Write.Length;
	buffer = irp->AssociatedIrp.SystemBuffer;

	device_extension->buffer  =	ExAllocatePool(PagedPool, size);

	if (device_extension->buffer  == NULL) {
		status = STATUS_INSUFFICIENT_RESOURCES;
		size = 0;
	} else {
		device_extension->buffer_size = size;
		RtlCopyMemory(device_extension->buffer, buffer,size);
	}

	irp->IoStatus.Status = status;
	irp->IoStatus.Information = size;
	IoCompleteRequest(irp, IO_NO_INCREMENT);
	return status;
}

static NTSTATUS Read (PDEVICE_OBJECT device, PIRP irp) 
{	
	PIO_STACK_LOCATION irpStack = IoGetCurrentIrpStackLocation(irp);
	PDEVICE_EXTENSION device_extension = (PDEVICE_EXTENSION) device->DeviceExtension;
	ULONG size = irpStack->Parameters.Read.Length;
	
	size = (size < device_extension->buffer_size) ? size : device_extension->buffer_size;

	RtlCopyMemory(irp->AssociatedIrp.SystemBuffer, device_extension->buffer, size);

	ExFreePool(device_extension->buffer);
	device_extension->buffer = NULL;
	device_extension->buffer_size = 0;

	irp->IoStatus.Status = STATUS_SUCCESS;
	irp->IoStatus.Information = size;
	IoCompleteRequest(irp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}

NTSTATUS DriverEntry (PDRIVER_OBJECT driver,PUNICODE_STRING pRegistryPath) 
{
	driver->DriverUnload = DriverUnload;
	driver->MajorFunction[IRP_MJ_CREATE] = Create;
	driver->MajorFunction[IRP_MJ_CLOSE]	 = Close;
	driver->MajorFunction[IRP_MJ_WRITE]	 = Write;
	driver->MajorFunction[IRP_MJ_READ]	 = Read;

	return CreateDevice(driver, 0);	
}

